#ifndef _AGRAD_HXX__
#define _AGRAD_HXX__

#include <vector>
#include <tuple>
#include <torch/extension.h>

// torch::Tensor CustomLinear(torch::Tensor input, torch::Tensor weight, torch::optional<torch::Tensor> bias);
// torch::Tensor CustomDualLinear(torch::Tensor, torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>, torch::optional<torch::Tensor>);
// torch::Tensor CustomGCN(torch::Tensor input, 
//                         torch::Tensor row_offset, torch::Tensor adj_list, torch::optional<torch::Tensor> edge_weight,
//                         torch::Tensor r_row_offset, torch::Tensor r_adj_list, torch::optional<torch::Tensor> r_edge_weight);

namespace customgcn{
  // temporary inlet
  // void init_layer(const std::string);
  void init_graph_cpp(bool, torch::Tensor, torch::optional<torch::Tensor>);
  torch::Tensor gcn_norm(torch::Tensor row_offset, torch::Tensor adj_list, 
                          torch::Tensor odegree, torch::Tensor edge_data);
  // int64_t get_num_vertex();:
  torch::Tensor propagate(int64_t opcode,
                          torch::Tensor row_offset, 
                          torch::Tensor adj_list, 
                          torch::optional<torch::Tensor> row_index,
                          torch::optional<torch::Tensor> edge_data,
                          torch::optional<torch::Tensor> in_feat);

  torch::Tensor propagate_backward(torch::Tensor input, 
                                   torch::Tensor sorted_index, 
                                   torch::Tensor row_offset, 
                                   torch::Tensor adj_list, 
                                   torch::optional<torch::Tensor> edge_weight);
  
  std::tuple<torch::Tensor, torch::Tensor>
  fused_linear_propagate(
    torch::Tensor input,
    torch::Tensor weight,
    torch::optional<torch::Tensor> bias,
    torch::Tensor row_offset,
    torch::Tensor adj_list,
    torch::optional<torch::Tensor> edge_data);

  std::tuple<torch::Tensor, torch::Tensor>
  fused_linear_propagate_reversed(
    torch::Tensor input,
    torch::Tensor weight,
    torch::Tensor row_offset,
    torch::Tensor adj_list,
    torch::optional<torch::Tensor> edge_data);

  // torch::Tensor
  // propagate_chunked(
  //   torch::optional<torch::Tensor> bos, torch::optional<torch::Tensor> boe, torch::optional<torch::Tensor> bosb, torch::optional<torch::Tensor> hte, torch::optional<torch::Tensor> htw, 
  //   torch::optional<torch::Tensor> bosbd, torch::optional<torch::Tensor> dmd, 
  //   torch::optional<torch::Tensor> boel, torch::optional<torch::Tensor> lsl, torch::optional<torch::Tensor> ldps, torch::optional<torch::Tensor> lwl, 
  //   torch::Tensor bodl, torch::Tensor ad,
  //   int64_t dcnt, int64_t hcnt, int64_t lcnt, int64_t sbsize,
  //   torch::Tensor src_tensor);

  // torch::Tensor
  // propagate_gcoo(
  //   torch::Tensor input, torch::Tensor el, torch::Tensor go, torch::optional<torch::Tensor> wl);
};

namespace customgat {
  torch::Tensor apply_edge(int64_t opcode,
                           torch::Tensor row_offset,
                           torch::Tensor col_indx,
                           torch::optional<torch::Tensor> row_indx,
                           torch::optional<torch::Tensor> edge_weight,
                           torch::optional<torch::Tensor> feat_l,
                           torch::optional<torch::Tensor> feat_r);
};

namespace gcnhelper{
torch::Tensor gcn_norm_csr(
  torch::Tensor row_offset, 
  torch::Tensor adj_list, 
  torch::Tensor odegree, 
  torch::Tensor weight_list);


torch::Tensor gcn_norm_coo(
  torch::Tensor edge_list,
  torch::Tensor weight_list,
  const int64_t nvertex);
}

void init_stream();

#endif